import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# 参数设置
N = 8  # 天线数量
d = 0.5  # 天线间距
lambda_ = 1.0  # 波长
num_angles = 37  # 角度范围
theta_range = torch.linspace(-90, 90, num_angles)  # 角度范围（度）
theta_rad = theta_range * np.pi / 180  # 角度转换为弧度

# 定义神经网络
class BeamformerNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(BeamformerNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化网络、损失函数和优化器
input_size = 2  # sin(theta) 和 cos(theta)
hidden_size = 16
output_size = N  # 权重向量的长度
net = BeamformerNet(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(net.parameters(), lr=0.01)

# 训练模型
best_loss = float('inf')
for epoch in range(1000):
    optimizer.zero_grad()
    # 输入数据：sin(theta) 和 cos(theta)
    sin_theta = torch.sin(theta_rad).unsqueeze(1)  # (num_angles, 1)
    cos_theta = torch.cos(theta_rad).unsqueeze(1)  # (num_angles, 1)
    input = torch.cat([sin_theta, cos_theta], dim=1)  # (num_angles, 2)
    # 前向传播
    w = net(input)  # (num_angles, N)
    w = w / torch.norm(w, dim=1, keepdim=True)  # 归一化权重
    # 计算相位
    n = torch.arange(N).float().unsqueeze(0)  # (1, N)
    phi = 2 * np.pi * d * n * torch.sin(theta_rad.unsqueeze(1)) / lambda_  # (num_angles, N)
    # 计算波束形成输出
    Y = torch.abs(torch.sum(w * torch.exp(1j * phi), dim=1))**2  # (num_angles,)
    # 计算目标角度的功率
    target_power = Y.gather(0, torch.arange(num_angles).long())
    # 计算旁瓣功率
    sidelobe_mask = torch.abs(theta_range.unsqueeze(0) - theta_range.unsqueeze(1)) > 30  # (num_angles, num_angles)
    sidelobe_power = torch.max(Y.unsqueeze(0) * sidelobe_mask.float(), dim=1)[0]  # (num_angles,)
    # 计算损失：最大化主瓣功率，最小化旁瓣功率
    loss = -(torch.mean(target_power) - torch.mean(sidelobe_power))
    # 反向传播和优化
    loss.backward()
    optimizer.step()
    # 打印损失
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')
    # 保存最佳模型
    if loss.item() < best_loss:
        best_loss = loss.item()
        torch.save(net.state_dict(), 'best_beamformer.pth')

# 评估模型
net.load_state_dict(torch.load('best_beamformer.pth'))
net.eval()

# 绘制所有角度的波束图
with torch.no_grad():
    for angle_idx in range(num_angles):
        theta_input = theta_rad[angle_idx].unsqueeze(0).unsqueeze(1)  # (1, 1)
        sin_theta = torch.sin(theta_input)
        cos_theta = torch.cos(theta_input)
        input = torch.cat([sin_theta, cos_theta], dim=1)
        w = net(input)
        w = w / torch.norm(w, dim=1, keepdim=True)
        phi = 2 * np.pi * d * torch.arange(N).float().unsqueeze(0) * torch.sin(theta_rad.unsqueeze(1)) / lambda_
        Y = torch.abs(torch.sum(w.unsqueeze(1) * torch.exp(1j * phi), dim=2))**2
        Y = Y.squeeze().numpy()
        plt.plot(theta_range.numpy(), Y, label=f'Target Angle: {theta_range[angle_idx]}')
    plt.xlabel('Angle (degrees)')
    plt.ylabel('Power')
    plt.title('Beam Patterns for All Angles')
    plt.legend()
    plt.grid(True)
    plt.show()

# 计算波束宽度和旁瓣水平
def calculate_beamwidth_and_sidelobe(Y, target_angle_idx):
    Y_softmax = torch.softmax(Y, dim=1)
    target_prob = Y_softmax.gather(1, target_angle_idx.unsqueeze(1)).squeeze()
    # 计算主瓣宽度约为60度
    beamwidth = 60
    # 计算旁瓣水平
    mainlobe_mask = torch.abs(theta_range.unsqueeze(0) - theta_range[target_angle_idx].unsqueeze(1)) <= 30
    sidelobe_prob = torch.max(Y_softmax * (~mainlobe_mask).float(), dim=1)[0]
    sidelobe_level = sidelobe_prob / target_prob
    return beamwidth, sidelobe_level

# 评估模型并计算指标
with torch.no_grad():
    beamwidths = []
    sidelobe_levels = []
    for angle_idx in range(num_angles):
        theta_input = theta_rad[angle_idx].unsqueeze(0).unsqueeze(1)  # (1, 1)
        sin_theta = torch.sin(theta_input)
        cos_theta = torch.cos(theta_input)
        input = torch.cat([sin_theta, cos_theta], dim=1)
        w = net(input)
        w = w / torch.norm(w, dim=1, keepdim=True)
        phi = 2 * np.pi * d * torch.arange(N).float().unsqueeze(0) * torch.sin(theta_rad.unsqueeze(1)) / lambda_
        Y = torch.abs(torch.sum(w.unsqueeze(1) * torch.exp(1j * phi), dim=2))**2
        Y = Y.squeeze()  # (num_angles,)
        beamwidth, sidelobe_level = calculate_beamwidth_and_sidelobe(Y.unsqueeze(0), torch.tensor([angle_idx]))
        beamwidths.append(beamwidth)
        sidelobe_levels.append(sidelobe_level.item())
    # 绘制指标
    plt.subplot(2, 1, 1)
    plt.plot(theta_range.numpy(), beamwidths)
    plt.xlabel('Angle (degrees)')
    plt.ylabel('Beamwidth (degrees)')
    plt.title('Beamwidth for All Angles')
    plt.grid(True)
    plt.subplot(2, 1, 2)
    plt.plot(theta_range.numpy(), sidelobe_levels)
    plt.xlabel('Angle (degrees)')
    plt.ylabel('Sidelobe Level')
    plt.title('Sidelobe Level for All Angles')
    plt.grid(True)
    plt.show()
